package websocket

import (
	"github.com/gorilla/websocket"
	"sync"
)

type ConnectionManager struct {
	clients map[*websocket.Conn]bool
	mu      sync.Mutex
}

func NewConnectionManager() *ConnectionManager {
	return &ConnectionManager{
		clients: make(map[*websocket.Conn]bool),
	}
}

func (cm *ConnectionManager) AddClient(client *websocket.Conn) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.clients[client] = true
}

func (cm *ConnectionManager) RemoveClient(client *websocket.Conn) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	delete(cm.clients, client)
}

func (cm *ConnectionManager) Broadcast(message []byte) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	for client := range cm.clients {
		if err := client.WriteMessage(websocket.TextMessage, message); err != nil {
			client.Close()
			delete(cm.clients, client)
		}
	}
}
